# Aliyun EMR Serverless Spark Notebook Source
# Version v1

# Cell markdown
# # <center>本地部署开源大模型
# Cell markdown
# # 3.手动实现基于vLLM框架的chat()方法
# Cell markdown
# &emsp;&emsp;所谓续写模式和对话模式，源于模型的训练数据不同，对于Qwen-7B基座模型来说，其训练数据的格式是这样的：
# ```markdown
# 输入：1月30日上午，中央气象台官微发布一张郑州市降水预报图，直呼：“河南郑州本次降雪预报图，这么离谱的预报图头一次见！<|endoftext|>
# 输出：预计1月31日至2月5日，我国中东部地区将进入入冬以来持续时间最长、影响范围最广的雨雪冰冻天气过程。10省份将出现暴雪或大暴雪，河南、河北、山东、辽宁、湖北等地的日降水量或累计降雪量具有极端性，河南、湖北、安徽、湖南、贵州5省将出现冻雨。<|endoftext|>
# ...
# ...
# 输入：1月28日，“格力2024全球梦想盛典”在珠海举行。格力电器员工在“格力全球梦想盛典”上合唱《我妈董明珠》的视频引起网友热议。<|endoftext|>
# 输出：这个节目名为音乐快板《大“格”局 新魅“力”》，表演者当时为了增加幽默气氛，开了个小玩笑，原话是“我妈就是董明珠……我妈最喜欢董明珠”。<|endoftext|>
# ```
# Cell markdown
# &emsp;&emsp;而对话模型，其训练数据往往是一问一答的形式。每个模型都有其开发者自定义的输入输出格式，比如gpt3.5是这样的：
# ```python
# {"messages": [{"role": "system", "content": "Marv is a factual chatbot that is also sarcastic."}, 
#               {"role": "user", "content": "What's the capital of France?"}, 
#               {"role": "assistant", "content": "Paris, as if everyone doesn't know that already."}]}
# 
# {"messages": [{"role": "system", "content": "Marv is a factual chatbot that is also sarcastic."}, 
#               {"role": "user", "content": "Who wrote 'Romeo and Juliet'?"}, 
#               {"role": "assistant", "content": "Oh, just some guy named William Shakespeare. Ever heard of him?"}]}
# 
# {"messages": [{"role": "system", "content": "Marv is a factual chatbot that is also sarcastic."}, 
#               {"role": "user", "content": "How far is the Moon from Earth?"}, 
#               {"role": "assistant", "content": "Around 384,400 kilometers. Give or take a few, like that really matters."}]},]}]}
# ```
# Cell markdown
# &emsp;&emsp;Qwen微调的对话格式是这样的：
# ```python
# [
#   {
#     "id": "identity_0",
#     "conversations": [
#       {
#         "from": "user",
#         "value": "你好"
#       },
#       {
#         "from": "assistant",
#         "value": "我是一个语言模型，我叫通义千问。"
#       }
#     ]
#   }
# ]
# ```
# Cell markdown
# &emsp;&emsp;Qwen的开源系列中的Chat类对话模型（Qwen-7B-Chat），它是基于对应的预训练版本（Qwen-7B）进行微调（SFT，S监督学习，FT微调）得到的，所以基于上述对话形式，就会通过如下方式，来界定某一个角色下的输入和输出：
# ```json
# 输入：<|im_start|>system
#       你是一个智能助手。<|im_end|>                 （这一部分是系统信息）
#       
#       <|im_start|>user
#       你好呀，请你介绍一下你自己<|im_end|>
#       <|im_start|>assitant
#       我是Qwen-Chat模型，一个非常厉害的人工智能助手，你可以向我提出任意的提问<|im_end|>
#       <|im_start|>user
#       好的，我想吃点甜食，你能不能帮我推荐一些好吃的食物？<|im_end|>
#       <|im_start|>assitant
#       当然可以，如果你喜欢吃甜食，那么推荐你吃巧克力的甜甜圈，太好吃了<|im_end|>       （这一部分是历史对话信息）
#       
#       <|im_start|>user
#       那么去哪里可以买到呢？<|im_end|>
#        <|im_start|>assitant<|endoftext|>    （这一部分是最新一轮的用户提问）
# 
# 输出：你可以去超市、蛋糕店等地方买到这个美味的食物<|im_end|><|endoftext|>     （这部分是Qwen模型针对最新一轮的提问，结合历史内容生成的回复）
# ```
# Cell markdown
# &emsp;&emsp;所以我们能表象看到的微调格式，并不是实际意义上送入Qwen模型训练的形式，而是会经过一系列格式的转化后，再送入模型。可以在这里看到：
# Cell markdown
# <div align=center><img src="https://snowball101.oss-cn-beijing.aliyuncs.com/img/202402011611889.png" width=100%></div>
# Cell markdown
# &emsp;&emsp;这种格式，被Qwen官方定义为`ChatML`格式。上述的输出，对应的就是我们上一步使用chat()方法，对Qwen模型执行的三次对话，也就是这样的三次操作：
# ```python
# # 第一轮对话
# response, history = model.chat(tokenizer, "你好呀，请你介绍一下你自己", history="你是一个智能助手。")
# print(response)   -- >  我是Qwen-Chat模型，一个非常厉害的人工智能助手，你可以向我提出任意的提问
# 
# # 第一轮对话
# response, history = model.chat(tokenizer, "好的，我想吃点甜食，你能不能帮我推荐一些好吃的食物？", history=history)
# print(response)  -- >  当然可以，如果你喜欢吃甜食，那么推荐你吃巧克力的甜甜圈，太好吃了
# 
# # 第三轮对话
# response, history = model.chat(tokenizer, "那么去哪里可以买到呢？", history=history)
# print(response)  -- > 你可以去超市、蛋糕店等地方买到这个美味的食物
# 
# ```
# Cell markdown
# &emsp;&emsp;这里的<|im_end|>、<|im_start|>及<|endoftext|>，也是Qwen官方定义的特殊Token，其中Qwen-7B中有<|endoftext|>，Qwen-7B-Chat中有<|endoftext|>、<|im_start|>以及<|im_end|>。 同时词表中也留有供扩展的特殊token位，可用<|extra_0|>到<|extra_204|>来指代。我们可以通过`tokenizer.special_tokens`来查看： 
# Cell code
## 导入Transformer依赖库
from transformers import AutoModelForCausalLM, AutoTokenizer

## 初始化tokenizer
tokenizer = AutoTokenizer.from_pretrained("/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-7B-Chat/", trust_remote_code=True,)

## 打印特殊Token
tokenizer.special_tokens
# Cell code
tokenizer.special_tokens
# Cell markdown
# &emsp;&emsp;了解了对话格式和特殊Token标识后，我们能明确的就是应该把输入的自然语言形式的Prompt处理成哪种形式，其实也就是`ChatML`格式，然后送入Qwen模型执行推理过程，那我们就可以先把大的框架搭建起来，一步一步的实现这个过程。
# Cell markdown
# &emsp;&emsp;首先定义一个类，其核心就是__init__方法和chat方法，前者用来初始化模型和配置，后者是实现对话交互的主函数。
# Cell code
class myvLLMWrapper:
    def __init__(self):
        """
        :param
        :param
        :return
        :description 模型初始化配置
        """
        pass


    def chat(self, prompt, history=None,system="你是一个乐于助人的智能机器人。"):
        """
        :param 
        :param 
        :return: 
        :description 大模型交互主逻辑
        """
        pass
# Cell markdown
# &emsp;&emsp;第一步：加载大模型的三个常规操作：model、tokenizer和config，这部分在__init__函数中写：
# Cell code
## 导入Transformer依赖库
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig
from vllm import LLM


class myvLLMWrapper:
    def __init__(self, model_dir):
        """
        :param model_dir ： 大模型实际存储位置
        :param
        :return: Description of return value
        """

        # 加载配置，即Qweb模型目录下的generation_config.json文件，是推理的关键参数
        self.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

        # 使用vLLM框架加载模型
        self.model = LLM(model=model_dir, tokenizer=model_dir, trust_remote_code=True)

    def chat(self, prompt, history=None, system="You are a helpful assistant.", extra_stop_words_ids=[]):
        """
        :param
        :param
        :return: Description of return value
        """
        pass
# Cell markdown
# &emsp;&emsp;完成模型启动的配置代码后，先测试一下该流程是否能够正常加载模型：
# Cell markdown
# &emsp;&emsp;为了加速测试流程，我这里下载了qwen1.8b的int模型，下面这段代码可直接下载，大家也可以尝试使用小模型来测试代码流程，提升效率。
# Cell code
## from modelscope.hub.snapshot_download import snapshot_download

## model_dir = snapshot_download('qwen/Qwen-1_8B-Chat-Int4',cache_dir='./models',revision='master')
# Cell markdown
# > 如果报错：AssertionError: tensor model parallel group is already initialized， 原因是前面的代码已经加载模型，需要重启Notebook再次执行，释放资源后重新加载模型
# Cell code
model_dir = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat-Int4/"

vllm_model = myvLLMWrapper(model_dir)
# Cell markdown
# &emsp;&emsp;可以看到，Qwen-7B的模型，此时已经可以通过vLLM框架来托管启动。那么接下来，需要做的就是如何与已经启动的Qwen-7B-Chat模型来进行对话交互，即实现Chat()方法。
# Cell markdown
# &emsp;&emsp;首先，需要进行Pompt的预处理工作，也就是将`你好呀，请你介绍一下你自己` 、 `好的，我想吃点甜食，你能不能帮我推荐一些好吃的食物？` 这种自然语言的Prompt，格式化成Qwen期望输入的`ChatML`类型的对话格式。即：
# ```json
# 输入：<|im_start|>system
#       你是一个智能助手。<|im_end|>                 （这一部分是系统信息）
#       
#       <|im_start|>user
#       你好呀，请你介绍一下你自己<|im_end|>
#       <|im_start|>assitant
#       我是Qwen-Chat模型，一个非常厉害的人工智能助手，你可以向我提出任意的提问<|im_end|>
#       <|im_start|>user
#       好的，我想吃点甜食，你能不能帮我推荐一些好吃的食物？<|im_end|>
#       <|im_start|>assitant
#       当然可以，如果你喜欢吃甜食，那么推荐你吃巧克力的甜甜圈，太好吃了<|im_end|>       （这一部分是历史对话信息）
#       
#       <|im_start|>user
#       那么去哪里可以买到呢？<|im_end|>
#        <|im_start|>assitant<|endoftext|>    （这一部分是最新一轮的用户提问）
# 
# 输出：你可以去超市、蛋糕店等地方买到这个美味的食物<|im_end|><|endoftext|>     （这不部分是Qwen模型针对最新一轮的提问，结合历史内容生成的回复）
# ```
# Cell markdown
# &emsp;&emsp;在数据格式预处理过程中，我们可以分步来操作，首先来处理system部分，即：
# ```json
# <|im_start|>system
#       你是一个智能助手。<|im_end|> 
# ```
# Cell markdown
# &emsp;&emsp;这里有几个关键的Token，分别是：<|im_start|>、<|im_end|> 以及隐藏的 `\n`。 可以通过tokenizer.encode()方法，找到它们对应的input_ids：
# Cell code
## 初始化tokenizer
tokenizer = AutoTokenizer.from_pretrained("/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-7B-Chat/", trust_remote_code=True,)
# Cell code
tokenizer.encode("\n")
# Cell code
tokenizer.encode('system')
# Cell code
tokenizer.encode('user')
# Cell code
tokenizer.encode('你是一个智能助手。')
# Cell code
role='system'
prompt = '你是一个智能助手'
# Cell code
print(f"{role}\n{prompt}",tokenizer.encode('system') + tokenizer.encode('\n') + tokenizer.encode('你是一个智能助手。'))
# Cell markdown
# &emsp;&emsp;根据上述测试，我们可以定义一个函数，用来将 输入的 role 和 prompt 拼接起来。
# Cell code
def _format_tokenize(role, prompt):
    return f"{role}\n{prompt}",tokenizer.encode(role) + tokenizer.encode('\n') + tokenizer.encode(prompt)
# Cell markdown
# 测试一下：
# Cell code
system_text, system_tokens = _format_tokenize('system', '你是一个智能助手')
# Cell code
system_text
# Cell code
system_tokens
# Cell markdown
# 加入特殊字符<|im_start|>和<|im_end|>
# Cell code
im_start,im_start_tokens='<|im_start|>',[tokenizer.im_start_id]
im_end,im_end_tokens='<|im_end|>',[tokenizer.im_end_id]
# Cell code
system_full = f'{im_start}{system_text}{im_end}'
system_full
# Cell code
system_tokens_full = im_start_tokens + system_tokens + im_end_tokens
system_tokens_full
# Cell markdown
# &emsp;&emsp;第一部分测试完成，我们把这部分的代码逻辑写进myvLLMWrapper类中，并测试一下system部分的生成情况。
# Cell markdown
# &emsp;&emsp;这里在myvLLMWrapper类中新定义一个format_prompt方法，用来处理这个逻辑。
# Cell code
## 导入Transformer依赖库
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig
from vllm import LLM


class myvLLMWrapper:
    def __init__(self, model_dir):
        """
        :param model_dir ： 大模型实际存储位置
        :param
        :return: Description of return value
        """

        # 加载配置，即Qweb模型目录下的generation_config.json文件，是推理的关键参数
        self.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

        # 使用vLLM框架加载模型
        self.model = LLM(model=model_dir, tokenizer=model_dir, trust_remote_code=True)

        # 加入特殊Token
        self.tokenizer.eos_token_id=self.generation_config.eos_token_id
        # 推理终止词，遇到这些词停止继续推理
        self.stop_words_ids=[self.tokenizer.im_start_id,self.tokenizer.im_end_id,self.tokenizer.eos_token_id]
    
        # 可用的Tokens数量，用来控制单次最大输入Tokens
        self.token_space=self.generation_config.max_window_size


    def format_prompt(self, tokenizer, prompt, history=None, system=""):
        
        im_start, im_start_tokens = '<|im_start|>', [tokenizer.im_start_id]
        im_end, im_end_tokens = '<|im_end|>', [tokenizer.im_end_id]

        def _format_tokenize(role, prompt):
            """
            :param role:system
            :param prompt:你是一个智能助手
            :return:('{role}\n{content}', [8948, 198, 56568, 101909, 100168, 110498])
            """
            return f"{role}\n{prompt}", tokenizer.encode(role) + tokenizer.encode('\n') + tokenizer.encode(prompt)


        # System 角色的处理逻辑
        system_text, system_tokens = _format_tokenize('system', system)
        # 假设系统角色输入为：You are a helpful assistant， 经过这一步，会变成  <|im_start|>system
        #                                                                            You are a helpful assistant.<|im_end|>
        system_text_full = f'{im_start}{system_text}{im_end}'

        # 这里是其对应的编码：[151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645]
        system_tokens_full = im_start_tokens + system_tokens + im_end_tokens
        self.token_space -= len(system_tokens_full)
        
        print(system_text_full) 
        print(system_tokens_full)
    

    def chat(self, prompt, history=None, system="You are a helpful assistant."):
        """
        :param
        :param
        :return: Description of return value
        """
        prompt_text, prompt_tokens = self.format_prompt(self.tokenizer, prompt, system=system)
# Cell code
model_dir = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat/"

vllm_model = myvLLMWrapper(model_dir)

vllm_model.chat(prompt="你好，请你介绍一下你自己", system='你是一个智能助手。')
# Cell markdown
# 忽略报错，这里先测试system_text_full和system_tokens_full是否可以正常输出，可以看到格式完全是匹配的。
# Cell markdown
# 接下来我们来看第二部分：
# ```json
#       <|im_start|>user
#       你好呀，请你介绍一下你自己<|im_end|>
#       <|im_start|>assitant   
# ```
# Cell code
def _format_tokenize(role, prompt):
    """
    :param role:system
    :param prompt:你是一个智能助手
    :return:('{role}\n{content}', [8948, 198, 56568, 101909, 100168, 110498])
    """
    return f"{role}\n{prompt}", tokenizer.encode(role) + tokenizer.encode('\n') + tokenizer.encode(prompt)
# Cell code
from transformers import AutoTokenizer

## 初始化tokenizer
tokenizer = AutoTokenizer.from_pretrained("/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-7B-Chat/", trust_remote_code=True,)

user_text, user_tokens = _format_tokenize('user', "你好呀，请你介绍一下你自己")
print(user_text, user_tokens)
# Cell code
im_start, im_start_tokens = '<|im_start|>', [tokenizer.im_start_id]
im_end, im_end_tokens = '<|im_end|>', [tokenizer.im_end_id]

user_text_full = f"\n{im_start}{user_text}{im_end}\n{im_start}assistant\n"
user_text_full
# Cell code
user_tokens_prefix = tokenizer.encode("\n") + im_start_tokens
print(user_tokens_prefix)
user_tokens_suffix = im_end_tokens + tokenizer.encode("\n") + im_start_tokens + tokenizer.encode('assistant') + tokenizer.encode("\n")
print(user_tokens_suffix)
# Cell code
user_tokens_full = user_tokens_prefix + user_tokens + user_tokens_suffix

print(user_tokens_full)
# Cell code
## 导入Transformer依赖库
from vllm import LLM
from vllm.sampling_params import SamplingParams
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig



class myvLLMWrapper:
    def __init__(self, model_dir):
        """
        :param model_dir ： 大模型实际存储位置
        :param
        :return: Description of return value
        """

        # 加载配置，即Qweb模型目录下的generation_config.json文件，是推理的关键参数
        self.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

        # 使用vLLM框架加载模型
        self.model = LLM(model=model_dir, tokenizer=model_dir, trust_remote_code=True)

        # 加入特殊Token
        self.tokenizer.eos_token_id=self.generation_config.eos_token_id
        # 推理终止词，遇到这些词停止继续推理
        self.stop_words_ids=[self.tokenizer.im_start_id,self.tokenizer.im_end_id,self.tokenizer.eos_token_id]
    
        # 可用的Tokens数量，用来控制单次最大输入Tokens
        self.token_space=self.generation_config.max_window_size


    def format_prompt(self, tokenizer, prompt, history=None, system=""):
        
        im_start, im_start_tokens = '<|im_start|>', [tokenizer.im_start_id]
        im_end, im_end_tokens = '<|im_end|>', [tokenizer.im_end_id]

        def _format_tokenize(role, prompt):
            """
            :param role:system
            :param prompt:你是一个智能助手
            :return:('{role}\n{content}', [8948, 198, 56568, 101909, 100168, 110498])
            """
            return f"{role}\n{prompt}", tokenizer.encode(role) + tokenizer.encode('\n') + tokenizer.encode(prompt)


        # System 角色的处理逻辑
        system_text, system_tokens = _format_tokenize('system', system)
        # 假设系统角色输入为：You are a helpful assistant， 经过这一步，会变成  <|im_start|>system
        #                                                                            You are a helpful assistant.<|im_end|>
        system_text_full = f'{im_start}{system_text}{im_end}'

        # 这里是其对应的编码：[151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645]
        system_tokens_full = im_start_tokens + system_tokens + im_end_tokens
        self.token_space -= len(system_tokens_full)
        
        # User 角色的处理逻辑
        user_text, user_tokens = _format_tokenize('user', prompt)
        user_tokens_prefix = tokenizer.encode("\n") + im_start_tokens
        user_tokens_suffix = im_end_tokens + tokenizer.encode("\n") + im_start_tokens + tokenizer.encode('assistant') + tokenizer.encode("\n")

        # 判断一下，避免输入的Prompt过长，过长就截断
        if len(user_tokens_prefix) + len(user_tokens) + len(user_tokens_suffix) > self.token_space: 
            user_tokens_len= self.token_space - len(user_tokens_prefix) - len(user_tokens_suffix)
            user_tokens = user_tokens[:user_tokens_len]
            user_text = self.tokenizer.decode(user_tokens)
            
        user_text_full = f"\n{im_start}{user_text}{im_end}\n{im_start}assistant\n"
        user_tokens_full = user_tokens_prefix + user_tokens + user_tokens_suffix
        
        self.token_space -= len(user_tokens_full)


        prompt_str = f'{system_text_full}{user_text_full}'
        prompt_tokens = system_tokens_full + user_tokens_full
        return prompt_str,prompt_tokens

    def chat(self, prompt, history=None, system="You are a helpful assistant."):
        """
        :param
        :param
        :return: Description of return value
        """
        prompt_text, prompt_tokens = self.format_prompt(self.tokenizer, prompt, system=system)
        
        print(prompt_text)       
# Cell code
model_dir = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat/"

vllm_model = myvLLMWrapper(model_dir)

vllm_model.chat(prompt="你好，请你介绍一下你自己", system='你是一个智能助手。')
# Cell markdown
# 实现一条完整的对话格式后，接下来我们进行模型推理部分的代码，即开始编写chat()函数。
# Cell code
## 导入Transformer依赖库
from vllm import LLM
from vllm.sampling_params import SamplingParams
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig

class myvLLMWrapper:
    def __init__(self, model_dir):
        """
        :param model_dir ： 大模型实际存储位置
        :param
        :return: Description of return value
        """

        # 加载配置，即Qweb模型目录下的generation_config.json文件，是推理的关键参数
        self.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

        # 使用vLLM框架加载模型
        self.model = LLM(model=model_dir, tokenizer=model_dir, trust_remote_code=True)

        # 加入特殊Token
        self.tokenizer.eos_token_id=self.generation_config.eos_token_id
        # 推理终止词，遇到这些词停止继续推理
        self.stop_words_ids=[self.tokenizer.im_start_id,self.tokenizer.im_end_id,self.tokenizer.eos_token_id]
    
        # 可用的Tokens数量，用来控制单次最大输入Tokens
        self.token_space=self.generation_config.max_window_size


    def format_prompt(self, tokenizer, prompt, history=None, system=""):
        
        im_start, im_start_tokens = '<|im_start|>', [tokenizer.im_start_id]
        im_end, im_end_tokens = '<|im_end|>', [tokenizer.im_end_id]

        def _format_tokenize(role, prompt):
            """
            :param role:system
            :param prompt:你是一个智能助手
            :return:('{role}\n{content}', [8948, 198, 56568, 101909, 100168, 110498])
            """
            return f"{role}\n{prompt}", tokenizer.encode(role) + tokenizer.encode('\n') + tokenizer.encode(prompt)


        # System 角色的处理逻辑
        system_text, system_tokens = _format_tokenize('system', system)
        # 假设系统角色输入为：You are a helpful assistant， 经过这一步，会变成  <|im_start|>system
        #                                                                            You are a helpful assistant.<|im_end|>
        system_text_full = f'{im_start}{system_text}{im_end}'

        # 这里是其对应的编码：[151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645]
        system_tokens_full = im_start_tokens + system_tokens + im_end_tokens
        self.token_space -= len(system_tokens_full)
        
        # User 角色的处理逻辑
        user_text, user_tokens = _format_tokenize('user', prompt)
        user_tokens_prefix = tokenizer.encode("\n") + im_start_tokens
        user_tokens_suffix = im_end_tokens + tokenizer.encode("\n") + im_start_tokens + tokenizer.encode('assistant') + tokenizer.encode("\n")

        # 判断一下，避免输入的Prompt过长，过长就截断
        if len(user_tokens_prefix) + len(user_tokens) + len(user_tokens_suffix) > self.token_space: 
            user_tokens_len= self.token_space - len(user_tokens_prefix) - len(user_tokens_suffix)
            user_tokens = user_tokens[:user_tokens_len]
            user_text = self.tokenizer.decode(user_tokens)
            
        user_text_full = f"\n{im_start}{user_text}{im_end}\n{im_start}assistant\n"
        user_tokens_full = user_tokens_prefix + user_tokens + user_tokens_suffix
        
        self.token_space -= len(user_tokens_full)


        prompt_str = f'{system_text_full}{user_text_full}'
        prompt_tokens = system_tokens_full + user_tokens_full
        return prompt_str,prompt_tokens

    def chat(self, prompt, history=None, system="You are a helpful assistant."):
        """
        :param
        :param
        :return: Description of return value
        """
        prompt_text, prompt_tokens = self.format_prompt(self.tokenizer, prompt, system=system)
        
        print(prompt_text)

        # VLLM请求配置
        sampling_params=SamplingParams(stop_token_ids=self.stop_words_ids)
        
        # 调用VLLM执行推理
        vllm_outputs=self.model.generate(prompt_token_ids=[prompt_tokens],sampling_params=sampling_params)

        # 返回最终结果
        response= vllm_outputs[0].outputs[0].text
        
        print(response)
# Cell code
model_dir = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat/"

vllm_model = myvLLMWrapper(model_dir)

vllm_model.chat(prompt="你好，请你介绍一下你自己", system='你是一个智能助手。')
# Cell code
vllm_model.chat(prompt="能帮我推荐几个北京的美食吗", system='你是一个美食家')
# Cell markdown
# 但这里发现问题，在模型生成的内容中会被截断，主要原因在于vLLM请求配置，我们这里丰富一下参数，只需要加这样一个参数 max_tokens
# Cell markdown
# | 参数 | 描述 |
# |------|------|
# | `n` | 返回给定提示的输出序列数量。 |
# | `best_of` | 从提示生成的输出序列的数量。从这些`best_of`序列中，返回前`n`个序列。`best_of`必须大于或等于`n`。默认情况下，`best_of`设置为`n`。 |
# | `presence_penalty` | 根据生成的文本中迄今为止出现的新词汇来惩罚新词汇的浮点数。大于0的值鼓励模型使用新词汇，小于0的值鼓励模型重复词汇。 |
# | `frequency_penalty` | 根据生成的文本中迄今为止新词汇的频率来惩罚新词汇的浮点数。大于0的值鼓励模型使用新词汇，小于0的值鼓励模型重复词汇。 |
# | `repetition_penalty` | 根据新词汇是否出现在提示和迄今为止生成的文本中来惩罚新词汇的浮点数。大于1的值鼓励模型使用新词汇，小于1的值鼓励模型重复词汇。 |
# | `temperature` | 控制抽样随机性的浮点数。较低的值使模型更确定性，较高的值使模型更随机。零意味着贪婪抽样。 |
# | `top_p` | 控制考虑的顶级词汇的累积概率的浮点数。必须在(0, 1]之间。设置为1以考虑所有词汇。 |
# | `top_k` | 控制考虑的顶级词汇数量的整数。设置为-1以考虑所有词汇。 |
# | `min_p` | 表示相对于最有可能的词汇的概率的最小概率的浮点数。必须在[0, 1]之间。设置为0以禁用此功能。 |
# | `use_beam_search` | 是否使用束搜索而不是抽样。 |
# | `length_penalty` | 根据序列的长度惩罚序列的浮点数。在束搜索中使用。 |
# | `early_stopping` | 控制束搜索的停止条件。可接受的值包括：`True`（生成停止，一旦有`best_of`个完整候选者）；`False`（应用启发式方法，当发现更好的候选者非常不可能时停止生成）；`"never"`（只有在不存在更好的候选者时束搜索过程才停止）。 |
# | `stop` | 停止生成时生成的字符串列表。返回的输出不会包含停止字符串。 |
# | `stop_token_ids` | 停止生成时生成的令牌列表。返回的输出将包含停止令牌，除非停止令牌是特殊令牌。 |
# | `include_stop_str_in_output` | 是否在输出文本中包含停止字符串。默认为False。 |
# | `ignore_eos` | 是否忽略EOS令牌，并在生成EOS令牌后继续生成令牌。 |
# | `max_tokens` | 每个输出序列生成的最大令牌数。 |
# | `logprobs` | 每个输出词汇返回的对数概率数量。 |
# | `prompt_logprobs` | 每个提示词汇返回的对数概率数量。 |
# | `skip_special_tokens` | 是否跳过输出中的特殊令牌。 |
# | `spaces_between_special_tokens` | 是否在输出中的特殊令牌之间添加空格。默认为
# Cell code
## 导入Transformer依赖库
from vllm import LLM
from vllm.sampling_params import SamplingParams
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig

class myvLLMWrapper:
    def __init__(self, model_dir):
        """
        :param model_dir ： 大模型实际存储位置
        :param
        :return: Description of return value
        """

        # 加载配置，即Qweb模型目录下的generation_config.json文件，是推理的关键参数
        self.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

        # 使用vLLM框架加载模型
        self.model = LLM(model=model_dir, tokenizer=model_dir, trust_remote_code=True)

        # 加入特殊Token
        self.tokenizer.eos_token_id=self.generation_config.eos_token_id
        # 推理终止词，遇到这些词停止继续推理
        self.stop_words_ids=[self.tokenizer.im_start_id,self.tokenizer.im_end_id,self.tokenizer.eos_token_id]
    
        # 可用的Tokens数量，用来控制单次最大输入Tokens
        self.token_space=self.generation_config.max_window_size


    def format_prompt(self, tokenizer, prompt, history=None, system=""):
        
        im_start, im_start_tokens = '<|im_start|>', [tokenizer.im_start_id]
        im_end, im_end_tokens = '<|im_end|>', [tokenizer.im_end_id]

        def _format_tokenize(role, prompt):
            """
            :param role:system
            :param prompt:你是一个智能助手
            :return:('{role}\n{content}', [8948, 198, 56568, 101909, 100168, 110498])
            """
            return f"{role}\n{prompt}", tokenizer.encode(role) + tokenizer.encode('\n') + tokenizer.encode(prompt)


        # System 角色的处理逻辑
        system_text, system_tokens = _format_tokenize('system', system)
        # 假设系统角色输入为：You are a helpful assistant， 经过这一步，会变成  <|im_start|>system
        #                                                                            You are a helpful assistant.<|im_end|>
        system_text_full = f'{im_start}{system_text}{im_end}'

        # 这里是其对应的编码：[151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645]
        system_tokens_full = im_start_tokens + system_tokens + im_end_tokens
        self.token_space -= len(system_tokens_full)
        
        # User 角色的处理逻辑
        user_text, user_tokens = _format_tokenize('user', prompt)
        user_tokens_prefix = tokenizer.encode("\n") + im_start_tokens
        user_tokens_suffix = im_end_tokens + tokenizer.encode("\n") + im_start_tokens + tokenizer.encode('assistant') + tokenizer.encode("\n")

        # 判断一下，避免输入的Prompt过长，过长就截断
        if len(user_tokens_prefix) + len(user_tokens) + len(user_tokens_suffix) > self.token_space: 
            user_tokens_len= self.token_space - len(user_tokens_prefix) - len(user_tokens_suffix)
            user_tokens = user_tokens[:user_tokens_len]
            user_text = self.tokenizer.decode(user_tokens)
            
        user_text_full = f"\n{im_start}{user_text}{im_end}\n{im_start}assistant\n"
        user_tokens_full = user_tokens_prefix + user_tokens + user_tokens_suffix
        
        self.token_space -= len(user_tokens_full)


        prompt_str = f'{system_text_full}{user_text_full}'
        prompt_tokens = system_tokens_full + user_tokens_full
        return prompt_str,prompt_tokens

    def chat(self, prompt, history=None, system="You are a helpful assistant."):
        """
        :param
        :param
        :return: Description of return value
        """
        prompt_text, prompt_tokens = self.format_prompt(self.tokenizer, prompt, system=system)
        
        print(prompt_text)

        # VLLM请求配置
        sampling_params=SamplingParams(stop_token_ids=self.stop_words_ids, max_tokens=1024)
        
        # 调用VLLM执行推理
        vllm_outputs=self.model.generate(prompt_token_ids=[prompt_tokens],sampling_params=sampling_params)

        # 返回最终结果
        response= vllm_outputs[0].outputs[0].text
        
        print(response)
# Cell code
model_dir = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat/"

## 实例化
vllm_model = myvLLMWrapper(model_dir)
# Cell code
vllm_model.chat(prompt="能帮我推荐几个北京的美食吗", system='你是一个美食家')
# Cell code
vllm_model.chat(prompt="我想去滑雪，我应该买什么滑雪设备", system='你是一个滑雪高手')
# Cell markdown
# &emsp;&emsp;最后，我们再丰富一下，添加上History信息。
# Cell code
## 导入Transformer依赖库
from vllm import LLM
from vllm.sampling_params import SamplingParams
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig
import copy
import os

## 通义千问的特殊token
IMSTART = '<|im_start|>'
IMEND = '<|im_end|>'
ENDOFTEXT = '<|endoftext|>'


class myvLLMWrapper:
    def __init__(self,
                 model_dir,):
        """
        :param model_dir ： 大模型实际存储位置
        :param
        :return: Description of return value
        """

        # 加载配置，即Qweb模型目录下的generation_config.json文件，是推理的关键参数
        self.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)

        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

        # 使用vLLM框架加载模型
        self.model = LLM(model=model_dir, tokenizer=model_dir, trust_remote_code=True)

        # 加入特殊Token
        self.tokenizer.eos_token_id = self.generation_config.eos_token_id
        # 推理终止词，遇到这些词停止继续推理
        self.stop_words_ids = [self.tokenizer.im_start_id, self.tokenizer.im_end_id, self.tokenizer.eos_token_id]

        # 可用的Tokens数量，用来控制单次最大输入Tokens
        self.token_space = self.generation_config.max_window_size

    def format_prompt(self,
                      tokenizer,
                      prompt,
                      history=None,
                      system=""):

        im_start, im_start_tokens = '<|im_start|>', [tokenizer.im_start_id]
        im_end, im_end_tokens = '<|im_end|>', [tokenizer.im_end_id]

        if history is None:
            history = []

        def _format_tokenize(role, prompt):
            """
            :param role:system
            :param prompt:你是一个智能助手
            :return:('{role}\n{content}', [8948, 198, 56568, 101909, 100168, 110498])
            """
            return f"{role}\n{prompt}", tokenizer.encode(role) + tokenizer.encode('\n') + tokenizer.encode(prompt)

        # System 角色的处理逻辑
        system_text, system_tokens = _format_tokenize('system', system)
        # 假设系统角色输入为：You are a helpful assistant， 经过这一步，会变成  <|im_start|>system
        #                                                                            You are a helpful assistant.<|im_end|>
        system_text_full = f'{im_start}{system_text}{im_end}'

        # 这里是其对应的编码：[151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645]
        system_tokens_full = im_start_tokens + system_tokens + im_end_tokens
        self.token_space -= len(system_tokens_full)

        # User 角色的处理逻辑
        user_text, user_tokens = _format_tokenize('user', prompt)
        user_tokens_prefix = tokenizer.encode("\n") + im_start_tokens
        user_tokens_suffix = im_end_tokens + tokenizer.encode("\n") + im_start_tokens + tokenizer.encode(
            'assistant') + tokenizer.encode("\n")

        # 判断一下，避免输入的Prompt过长，过长就截断
        if len(user_tokens_prefix) + len(user_tokens) + len(user_tokens_suffix) > self.token_space:
            user_tokens_len = self.token_space - len(user_tokens_prefix) - len(user_tokens_suffix)
            user_tokens = user_tokens[:user_tokens_len]
            user_text = self.tokenizer.decode(user_tokens)

        user_text_full = f"\n{im_start}{user_text}{im_end}\n{im_start}assistant\n"
        user_tokens_full = user_tokens_prefix + user_tokens + user_tokens_suffix

        self.token_space -= len(user_tokens_full)

        # 累积的历史对话文本和令牌列表
        accumulated_history_text, accumulated_history_tokens = '', []

        # 从最近的对话历史开始处理
        for previous_user_query, previous_assistant_response in reversed(history):
            # 格式化user用户的历史提问
            user_query_text, user_query_tokens = _format_tokenize("user", previous_user_query)
            # 格式化并分词assistant的历史回答
            assistant_response_text, assistant_response_tokens = _format_tokenize("assistant",
                                                                                  previous_assistant_response)

            # 构建当前轮次的对话token
            current_round_tokens = tokenizer.encode(
                "\n") + im_start_tokens + user_query_tokens + im_end_tokens + tokenizer.encode(
                "\n") + im_start_tokens + assistant_response_tokens + im_end_tokens
            # 构建当前轮次的对话文本
            current_round_text = f"\n{im_start}{user_query_text}{im_end}\n{im_start}{assistant_response_text}{im_end}"

            # 如果累积的对话令牌数量没有超过限制，则将本轮对话添加到累积历史中
            if len(current_round_tokens) <= self.token_space:
                accumulated_history_text = current_round_text + accumulated_history_text
                accumulated_history_tokens = current_round_tokens + accumulated_history_tokens
            else:
                break

        prompt_str = f'{system_text_full}{accumulated_history_text}{user_text_full}'
        prompt_tokens = system_tokens_full + accumulated_history_tokens + user_tokens_full
        return prompt_str, prompt_tokens

    def chat(self, prompt, history=None, system="You are a helpful assistant."):
        """
        :param
        :param
        :return: Description of return value
        """
        # 历史聊天
        if history is None:
            history = []
        else:
            history = copy.deepcopy(history)

        prompt_text, prompt_tokens = self.format_prompt(self.tokenizer, prompt, system=system, history=history)


        # VLLM请求配置
        sampling_params = SamplingParams(stop_token_ids=self.stop_words_ids,
                                         # 方式一：直接设置
                                         max_tokens=1024,
                                         # 方式二：去Qwen模型的`generation_config.json`文件中配置
                                         temperature=self.generation_config.temperature,)

        # 调用VLLM执行推理
        vllm_outputs=self.model.generate(prompt_token_ids=[prompt_tokens],sampling_params=sampling_params)

        # 返回最终结果
        response= vllm_outputs[0].outputs[0].text


        # 添加历史信息
        history.append((prompt, response))

        return response, history
# Cell code
model_dir = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat/"

## 实例化
vllm_model = myvLLMWrapper(model_dir)
# Cell code
response, history = vllm_model.chat(prompt="你好，请你帮我推荐几道北京的美食", system='你是一个美食家', history=None)
print(response)
# Cell code
response, history = vllm_model.chat(prompt="是吗？帮我具体展开聊一聊", system='你是一个美食家', history=history)
print(response)
# Cell code
response, history = vllm_model.chat(prompt="除了这些，还有什么呢？", system='你是一个美食家', history=history)
print(response)
# Cell markdown
# &emsp;&emsp;最后，我们基于Chat方法创建一个多轮对话，代码如下：
# Cell code
def initialize_model(model_path):
    """
    初始化并加载模型。
    """
    model = myvLLMWrapper(model_path)
    return model

def main():
    model_path = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat/"

    # 获取用户设置的历史记录长度
    try:
        history_length = int(input("请输入历史记录的长度（建议值：128）: "))
    except ValueError:
        print("输入无效，使用默认历史记录长度：128")
        history_length = 20

    history = None
    model_initialized = False

    while True:

        if not model_initialized:
            print("正在初始化模型，请稍等...")
            vllm_model = initialize_model(model_path)
            model_initialized = True
            print("模型初始化完成。")

        user_input = input('您的提问: ')
        if user_input.lower() == '退出':
            print("感谢您的使用，再见！")
            break


        try:
            response, history = vllm_model.chat(prompt=user_input, history=history)
            print("回答:", response)
            history = history[-history_length:]  # 使用用户设定的历史记录长度
        except Exception as e:
            print("发生错误:", e)
            break

if __name__ == "__main__":
    main()

# Cell markdown
# &emsp;&emsp;上述就可以看到，掌握了数据的输入，Transformers库对Qwen模型的加载方式，及明确模型的输出处理流程，一切中间的逻辑是完全可以自己手动去定义和封装的。除了处理的逻辑，像`myvLLMWrapper`的初始化，SamplingParams参数的设置，完全可以根据某个开源框架封装好的功能进行接入。比如：
# 
# ```python
# vllm_model = vLLMWrapper(model_path,
#                          dtype="float16",
#                          gpu_memory_utilization=0.6)
# ```
# Cell markdown
# > 详细的参数在这里：https://docs.vllm.ai/en/latest/genindex.html
# Cell markdown
# &emsp;&emsp;接入Gradio，生成网页端的Web UI。
# Cell code
import gradio as gr

## 模型和历史记录初始化
model_path = "/home/Work/00.Work_muyu/muyu_qwen/models/qwen/Qwen-1_8B-Chat/"  # 模型路径
vllm_model = myvLLMWrapper(model_path)

history = []

def chat_with_model(user_input):
    global history
    response, history = vllm_model.chat(prompt=user_input, history=history)
    history = history[-128:]  # 保持历史记录长度为128
    return response

iface = gr.Interface(
    fn=chat_with_model,
    inputs=gr.inputs.Textbox(lines=2, placeholder="输入您的问题..."),
    outputs="text",
    title="智能对话系统",
    description="输入您的问题并与智能模型进行对话。",
)

if __name__ == "__main__":
    iface.launch(server_name="192.168.110.131", server_port=12345)